import numpy as np
import matplotlib.pyplot as plt

# 生成随机数据集
np.random.seed(42)

# 数据集包含两个簇，每个簇包含100个样本
cluster1_mean = [2, 2]
cluster2_mean = [6, 6]
cluster1_cov = [[1, 0.5], [0.5, 1]]
cluster2_cov = [[1, -0.5], [-0.5, 1]]

cluster1 = np.random.multivariate_normal(cluster1_mean, cluster1_cov, 100)
cluster2 = np.random.multivariate_normal(cluster2_mean, cluster2_cov, 100)

data = np.vstack((cluster1, cluster2))

plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号 #有中文出现的情况，需要u'内容'

# 绘制数据集
plt.scatter(data[:, 0], data[:, 1], s=10)
plt.title('生成的数据')
plt.xlabel('X')
plt.ylabel('Y')
plt.show()


def k_means(data, k, max_iters=100):
    # 随机初始化聚类中心
    centroids = data[np.random.choice(data.shape[0], k, replace=False)]

    for _ in range(max_iters):
        # 计算每个样本到聚类中心的距离
        distances = np.sqrt(((data - centroids[:, np.newaxis])**2).sum(axis=2))

        # 分配样本到最近的聚类中心
        labels = np.argmin(distances, axis=0)

        # 更新聚类中心为每个簇的平均值
        for i in range(k):
            if np.any(labels == i):
                centroids[i] = np.mean(data[labels == i], axis=0)

    return centroids, labels

# 运行 K-means 算法
k = 2  # 要分成的簇数
centroids, labels = k_means(data, k)

# 绘制簇和聚类中心
plt.scatter(data[:, 0], data[:, 1], c=labels, s=10, cmap='viridis')
plt.scatter(centroids[:, 0], centroids[:, 1], marker='o', c='red', s=100, label='中心点')
plt.title('K-means聚类')
plt.xlabel('X')
plt.ylabel('Y')
plt.legend()
plt.show()